import torch
import os
from PIL import Image
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter

# 数据集获取部分
class ReadData(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classnames = os.listdir(root_dir)
        self.imgs, self.labels = self.get_images_labels()

    def get_images_labels(self):
        imgs = []
        labels = []

        for class_idx, class_name in enumerate(self.classnames):
            class_dir = os.path.join(self.root_dir, class_name)
            img_names = os.listdir(class_dir)
            for name in img_names:
                img_dir = os.path.join(class_dir, name)
                imgs.append(img_dir)
                labels.append(class_idx)

        return imgs, labels

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, idx):
        img = Image.open(self.imgs[idx]).convert("RGB")
        label = self.labels[idx]

        if self.transform:
            img = self.transform(img)

        return img, label

# 数据集处理部分
train_transform = transforms.Compose(
    [transforms.Resize((224, 224)),          # 调整图像大小
    transforms.RandomHorizontalFlip(),      # 随机水平翻转
    transforms.RandomRotation(15),          # 随机旋转（-15°到15°）
    transforms.ToTensor(),                  # 转换为 Tensor [0, 1]
    transforms.Normalize(                   # 标准化（ImageNet 均值和标准差）
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )]
)

test_transform = transforms.Compose(
    [transforms.Resize((224, 224)),          # 调整图像大小
    transforms.ToTensor(),                  # 转换为 Tensor [0, 1]
    transforms.Normalize(                   # 标准化（ImageNet 均值和标准差）
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )]
)

# 数据加载部分
train_dir = 'E:/代码仓库/测试数据集/train'
test_dir = 'E:/代码仓库/测试数据集/val'

train_dataset = ReadData(train_dir, train_transform)
test_dataset = ReadData(test_dir, test_transform)

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4
)

test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

def denormalize(tensor, mean, std):
    """反归一化函数"""
    mean = torch.tensor(mean).view(1, 3, 1, 1)
    std = torch.tensor(std).view(1, 3, 1, 1)
    denorm_tensor = tensor * std + mean
    return torch.clamp(denorm_tensor, 0, 1)  # 确保值在 [0, 1] 范围内


def test_DL(train_loader):
    log_dir = './run_log'
    writer = SummaryWriter(log_dir)
    idx = 1
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    for data in train_loader:
        imgs, labels = data
        # 反归一化图像
        imgs = denormalize(imgs.clone(), mean, std)
        writer.add_images(f"第{idx}批图像", imgs)
        print(imgs.size(), labels.size())
        idx += 1

    writer.close()

if __name__ == '__main__':
    test_DL(train_loader)